import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.models import resnet50


class FPN(nn.Module):
    """Feature Pyramid Network with ResNet50 backbone"""

    def __init__(self, num_classes=80):
        super(FPN, self).__init__()

        # 使用预训练的ResNet50作为backbone
        resnet = resnet50()

        # 获取各个stage的输出
        self.layer1 = nn.Sequential(
            resnet.conv1,
            resnet.bn1,
            resnet.relu,
            resnet.maxpool,
            resnet.layer1  # C2
        )
        self.layer2 = resnet.layer2  # C3
        self.layer3 = resnet.layer3  # C4
        self.layer4 = resnet.layer4  # C5

        # 横向连接层 (lateral connections)
        self.lateral_c4 = nn.Conv2d(1024, 256, 1)
        self.lateral_c3 = nn.Conv2d(512, 256, 1)
        self.lateral_c2 = nn.Conv2d(256, 256, 1)

        # 3x3卷积融合特征
        self.smooth_p4 = nn.Conv2d(256, 256, 3, padding=1)
        self.smooth_p3 = nn.Conv2d(256, 256, 3, padding=1)
        self.smooth_p2 = nn.Conv2d(256, 256, 3, padding=1)

        # 用于目标检测的头部网络
        self.detection_head = nn.Sequential(
            nn.Conv2d(256, 256, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, num_classes, 1)
        )

    def _upsample_add(self, x, y):
        """上采样x并与y相加"""
        _, _, H, W = y.shape
        # 论文中使用“nearest”
        return F.interpolate(x, size=(H, W), mode='bilinear', align_corners=True) + y

    def forward(self, x):
        # 自底向上的主干网络前向传播
        c2 = self.layer1(x)  # 1/4
        c3 = self.layer2(c2)  # 1/8
        c4 = self.layer3(c3)  # 1/16
        c5 = self.layer4(c4)  # 1/32

        # 横向连接
        p5 = nn.Conv2d(2048, 256, 1)(c5)
        p4 = self._upsample_add(p5, self.lateral_c4(c4))
        p3 = self._upsample_add(p4, self.lateral_c3(c3))
        p2 = self._upsample_add(p3, self.lateral_c2(c2))

        # 特征图平滑
        p4 = self.smooth_p4(p4)
        p3 = self.smooth_p3(p3)
        p2 = self.smooth_p2(p2)

        # 目标检测预测 【共享分类头】
        pred_p2 = self.detection_head(p2)
        pred_p3 = self.detection_head(p3)
        pred_p4 = self.detection_head(p4)
        pred_p5 = self.detection_head(p5)

        return [pred_p2, pred_p3, pred_p4, pred_p5]


# 使用示例
def demo():
    # 创建模型实例
    model = FPN(num_classes=80)

    # 生成示例输入
    batch_size = 1
    sample_input = torch.randn(batch_size, 3, 800, 800)

    # 前向传播
    predictions = model(sample_input)

    # 打印各个尺度的预测输出大小
    for i, pred in enumerate(predictions):
        print(f"P{i + 2} shape:", pred.shape)

    return model, predictions


# 测试代码
if __name__ == "__main__":
    model, preds = demo()
